package top.goldenyear.porpoise.common.core.util.tree;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import cn.hutool.v7.core.util.ObjUtil;

/**
 * @author airhead
 */
public class TreeUtils {

  public static <T, N extends BaseNode<T, N>> List<N> buildTree(List<N> sources) {
    Set<T> tempSet = sources.stream().map(BaseNode::getId).collect(Collectors.toSet());

    List<N> results = new ArrayList<>();
    // get root nodes
    List<N> rootNodes =
        sources.stream()
            .filter(x -> x.getParentId() == null || !tempSet.contains(x.getParentId()))
            .toList();

    for (N rootNode : rootNodes) {
      results.add(buildChildTree(sources, rootNode));
    }

    return results;
  }

  public static <T, N extends BaseNode<T, N>> List<N> buildTree(List<N> sources, T rootId) {
    List<N> results = new ArrayList<>();
    // get root nodes
    List<N> rootNodes =
        sources.stream()
            .filter(x -> x.getParentId() == null || ObjUtil.equals(x.getParentId(), rootId))
            .toList();
    for (N rootNode : rootNodes) {
      results.add(buildChildTree(sources, rootNode));
    }

    return results;
  }

  public static <N extends BaseNode<?, N>> N buildChildTree(List<N> sources, N pNode) {
    List<N> children = null;
    for (N source : sources) {
      if (source.getParentId() != null && source.getParentId().equals(pNode.getId())) {
        if (children == null) {
          children = new ArrayList<>();
        }

        children.add(buildChildTree(sources, source));
      }
    }

    pNode.setChildren(children);

    return pNode;
  }
}
